//
//  pipeline.cpp
//  yuzon
//
//  Created by Alvin on 2023/3/29.
//

#include "pipeline.hpp"
#include "utils.hpp"
#include "visualizer.hpp"

using namespace cv;
using namespace std;

Pipeline::Pipeline(const string &model_detection_path,
                   const string &label_list_path,
                   const string &model_keypoint_path,
                   const string &angle_list_path) :
draw_bbox_(kDrawBbox),
draw_pose_(kDrawPose),
bbox_line_thickness_(kBboxLineThickness),
pose_keypoint_num_(kPoseKeypointNum),
pose_circle_radius_(kPoseCircleRadius),
pose_line_thickness_(kPoseLineThickness),
state_(-2) {
  human_detector_.reset(new HumanDetector(model_detection_path,
                                          label_list_path));
  keypoint_detector_.reset(new KeypointDetector(model_keypoint_path));
  scorer_.reset(new Scorer(angle_list_path));
  smoother_.reset(new Smoother());
}

int Pipeline::Process(Mat &image, int frame) {
  cvtColor(image, image, COLOR_BGRA2RGB);
  
  human_results_.clear();
  keypoint_results_.clear();
  elapsed_times_.clear();
  
  int64_t t1 = GetCurrentTime();
  human_detector_->Predict(image, human_results_);
  elapsed_times_.emplace_back(GetElapsedTime(t1));
  
  int64_t t2 = GetCurrentTime();
  keypoint_detector_->Predict(image, human_results_, keypoint_results_);
  elapsed_times_.emplace_back(GetElapsedTime(t2));
  
  int64_t t3 = GetCurrentTime();
  smoother_->Process(keypoint_results_);
  elapsed_times_.emplace_back(GetElapsedTime(t3));
  
  int64_t t4 = GetCurrentTime();
  VisualizeResults(image, annotation_image_, human_results_, keypoint_results_,
                   draw_bbox_, draw_pose_, bbox_line_thickness_,
                   pose_keypoint_num_, pose_circle_radius_, pose_line_thickness_);
  elapsed_times_.emplace_back(GetElapsedTime(t4));
  
  int64_t t5 = GetCurrentTime();
  scorer_->Process(keypoint_results_, frame, state_, frame_scores_);
  elapsed_times_.emplace_back(GetElapsedTime(t5));
  
  return state_;
}

void Pipeline::Reset() {
  state_ = -2;
  smoother_->Reset();
  frame_scores_.clear();
}

void Pipeline::SetDrawBbox(bool draw_bbox) {
  draw_bbox_ = draw_bbox;
}

void Pipeline::SetDrawPose(bool draw_pose) {
  draw_pose_ = draw_pose;
}

void Pipeline::SetBboxLineThickness(int bbox_line_thickness) {
  bbox_line_thickness_ = bbox_line_thickness;
}

void Pipeline::SetPoseKeypointNum(int pose_keypoint_num) {
  if (pose_keypoint_num == 17 || pose_keypoint_num == 14)
    pose_keypoint_num_ = pose_keypoint_num;
}

void Pipeline::SetPoseCircleRadius(int pose_circle_radius) {
  pose_circle_radius_ = pose_circle_radius;
}

void Pipeline::SetPoseLineThickness(int pose_line_thickness) {
  pose_line_thickness_ = pose_line_thickness;
}

int Pipeline::GetScore() {
  return frame_scores_.empty() ? 0 : frame_scores_.back()[1];
}

int Pipeline::GetAverageScore() {
  if (frame_scores_.empty())
    return 0;
  
  int sum = 0;
  int num = 0;
  
  for (const vector<int> &frame_score : frame_scores_) {
    if (frame_score[1] > 0) {
      sum += frame_score[1];
      ++num;
    }
  }
  
  return round(sum / num);
}

vector<float> Pipeline::GetKeypoints() {
  if (keypoint_results_.empty())
    return vector<float>(pose_keypoint_num_);
  
  const KeypointResults &results = (pose_keypoint_num_ == 17 ?
                                    keypoint_results_ :
                                    ConvertKeypoints(keypoint_results_));
  
  return results[0].keypoints;
}

vector<float> Pipeline::GetTimes() {
  return elapsed_times_;
}

Mat Pipeline::GetAnnotationImage() {
  return annotation_image_;
}

vector<vector<int>> Pipeline::GetFrameScores() {
  return frame_scores_;
}
